// ECE 5984 S20 Assignment 2: available.cpp
// Group: Sunny Wadkar and Ramprasath Pitta Sekar

////////////////////////////////////////////////////////////////////////////////

#include "llvm/IR/Function.h"
#include "llvm/IR/PassManager.h" // 变更为新PM的头文件
#include "llvm/Passes/PassBuilder.h" // 新增：用于Pass构建
#include "llvm/Passes/PassPlugin.h" // 新增：用于插件支持
#include "llvm/IR/Constants.h"
#include "llvm/Support/raw_ostream.h"

#include "dataflow.h"
#include "available-support.h"

using namespace llvm;
using namespace std;

namespace {
  // 变更为继承PassInfoMixin，不再使用FunctionPass :cite[6]:cite[9]
  class AvailableExpressions : public PassInfoMixin<AvailableExpressions> {

  public:
    // 不再需要静态ID变量

    // 变更为新PM的run方法签名 :cite[6]
    PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM) {

      // Here's some code to familarize you with the Expression
      // class and pretty printing code we've provided:

      vector<Expression> expressions;
      for (Function::iterator FI = F.begin(), FE = F.end(); FI != FE; ++FI) {
               BasicBlock* block = &*FI;
               for (BasicBlock::iterator i = block->begin(), e = block->end(); i!=e; ++i) {

                 Instruction* I = &*i;
                 // We only care about available expressions for BinaryOperators
                 if (BinaryOperator *BI = dyn_cast<BinaryOperator>(I))
           {
             if(std::find(expressions.begin(),expressions.end(),Expression(BI)) != expressions.end())
             {
               continue;
             }
             else
             {
               expressions.push_back(Expression(BI));
             }
                 }
               }
      }

      // Print out the expressions used in the function
      outs() << "Expressions used by this function:\n";
      printSet(&expressions);

      BitVector boundaryCond(expressions.size(), false);
      BitVector initCond(expressions.size(), true);
      dataFlow* df = new dataFlow(expressions.size(),INTERSECTION,FORWARD,boundaryCond,initCond);
      genTFDependencies(F,expressions);
      df->executeDataFlowPass(F,bbSets);
      printDFAResults(df->dataFlowHash);
      
      // 返回值变更为PreservedAnalyses :cite[6]
      // 由于此Pass不修改IR，我们保留所有分析 :cite[6]
      return PreservedAnalyses::all();
    }

    // 不再需要getAnalysisUsage方法 :cite[1]

  private:
    std::map<BasicBlock*,basicBlockDeps*> bbSets;
    std::map<Expression,int> domainMap;
    std::map<int,string> revDomainMap;
    void genTFDependencies(Function &F, std::vector<Expression> domain)
    {
      BitVector empty((int)domain.size(), false);
      int vectorIdx = 0;
      StringRef lhs;
      for(Expression e : domain)
      {
        domainMap[e] = vectorIdx;
        revDomainMap[vectorIdx] = e.toString();
        ++vectorIdx;
      }
      for(BasicBlock &BB : F)
      {
        struct basicBlockDeps* bbSet = new basicBlockDeps();
        bbSet->blockRef = &BB;
        bbSet->genSet = empty;
        bbSet->killSet = empty;
        for(Instruction &II : BB)
        {
          Instruction *I = &II;
          if(dyn_cast<BinaryOperator>(I))
          {
            lhs = I->getName();
            if(std::find(domain.begin(),domain.end(),Expression(I)) != domain.end())
            {
              bbSet->genSet.set(domainMap[Expression(I)]);
            }
            for(Expression itr: domain)
            {
              StringRef exprOp1 = itr.v1->getName();
              StringRef exprOp2 = itr.v2->getName();
              if(exprOp1.equals(lhs) || exprOp2.equals(lhs))
              {
                bbSet->killSet.set(domainMap[itr]);
                bbSet->genSet.reset(domainMap[itr]);
              }
            }
          }
        }
        bbSets[&BB] = bbSet;
      }
    }

    void printDFAResults(std::map<BasicBlock*,basicBlockProps*> dFAHash)
    {
      std::map<BasicBlock*,basicBlockProps*>::iterator it = dFAHash.begin();
      while(it != dFAHash.end())
      {
        struct basicBlockProps* temp = dFAHash[it->first];
        outs() << "Basic Block Name : ";
        outs() << temp->block->getName() << "\n";
        std::vector<string> genbb;
        outs() << "gen[BB] : ";
        for(int m = 0; m < (int)temp->genSet.size(); m++)
        {
          if(temp->genSet[m])
          {
            genbb.push_back(revDomainMap[m]);
          }
        }
        printStringSet(&genbb);
        std::vector<string> killbb;
        outs() << "kill[BB] : ";
        for(int n = 0; n < (int)temp->killSet.size(); n++)
        {
          if(temp->killSet[n])
          {
            killbb.push_back(revDomainMap[n]);
          }
        }
        printStringSet(&killbb);
        std::vector<string> inbb;
        outs() << "IN[BB] : ";
        for(int k = 0; k < (int)temp->bbInput.size(); k++)
        {
          if(temp->bbInput[k])
          {
            inbb.push_back(revDomainMap[k]);
          }
        }
        printStringSet(&inbb);
        std::vector<string> outbb;
        outs() << "OUT[BB] : ";
        for(int l = 0; l < (int)temp->bbOutput.size(); l++)
        {
          if(temp->bbOutput[l])
          {
            outbb.push_back(revDomainMap[l]);
          }
        }
        printStringSet(&outbb);
        outs() << "\n";
        it++;
      }
    }
  };

  // 不再需要RegisterPass和静态ID注册 :cite[1]
} // 结束匿名命名空间

// 新增：新PM的插件注册接口 :cite[4]:cite[6]
extern "C" ::llvm::PassPluginLibraryInfo LLVM_ATTRIBUTE_WEAK
llvmGetPassPluginInfo() {
  return {
    LLVM_PLUGIN_API_VERSION, "AvailableExpressions", "v0.1",
    [](PassBuilder &PB) {
      PB.registerPipelineParsingCallback(
        [](StringRef Name, FunctionPassManager &FPM,
           ArrayRef<PassBuilder::PipelineElement>) {
          if (Name == "available") {
            FPM.addPass(AvailableExpressions());
            return true;
          }
          return false;
        });
    }
  };
}
